import numpy as np
import pickle
from elftools.elf.elffile import ELFFile
from elftools.elf.sections import SymbolTableSection
import math
import io

scaler = pickle.load(open("./static/model/elf/elf_scaler.pkl", "rb"))
model = pickle.load(open('./static/model/elf/elf_model.pkl', 'rb'))


def calculate_entropy(data):
    if not data:
        return 0
    entropy = 0
    for x in range(256):
        p_x = float(data.count(x)) / len(data)
        if p_x > 0:
            entropy += - p_x * math.log(p_x, 2)
    return entropy


def get_static(file):
    stream = io.BytesIO(file)
    elffile = ELFFile(stream)
    features = []
    # Basic ELF features
    features.append(elffile['e_shnum'])
    features.append(elffile['e_ehsize'])
    features.append(elffile['e_phentsize'])
    features.append(elffile['e_shentsize'])
    features.append(elffile['e_shstrndx'])

    # Section features
    sections = list(elffile.iter_sections())
    features.append(len(sections))
    entropies = [calculate_entropy(sec.data())
                 for sec in sections if sec.data()]
    raw_sizes = [len(sec.data()) for sec in sections if sec.data()]
    features.append(max(entropies, default=0))
    features.append(min(entropies, default=0))
    features.append(sum(entropies) / len(entropies) if entropies else 0)
    features.append(max(raw_sizes, default=0))
    features.append(min(raw_sizes, default=0))
    features.append(sum(raw_sizes) / len(raw_sizes) if raw_sizes else 0)

    # Symbol features
    symbols = []
    for section in sections:
        if isinstance(section, SymbolTableSection):
            symbols.extend(section.iter_symbols())
    symbol_sizes = [sym['st_size'] for sym in symbols]
    features.append(len(symbols))
    features.append(sum(symbol_sizes) / len(symbol_sizes)
                    if symbol_sizes else 0)
    features.append(max(symbol_sizes, default=0))
    features.append(min(symbol_sizes, default=0))

    return features


def model_predict(static):
    features_vec = np.array(static)
    vec = scaler.transform([features_vec])
    features_vec_np = np.array(vec)
    pred = model.predict(features_vec_np)
    return pred[0]


def process(file):
    static = get_static(file)
    return model_predict(static), static


def process_string(static):
    return model_predict(static), static


if __name__ == '__main__':
    # file_path = "../elf_samples/x86_64-linux-gnu-ld.bfd"
    file_path = "../elf_samples/VirusShare_1dbe0f7d4181fa0a5dac884757f0af95"
    with open(file_path, "rb") as f_in:
        file = f_in.read()
    res = process(file)
    print(res)
